import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.tensorboard import SummaryWriter

# 1. 生成随机数据
def generate_random_data(num_samples, input_dim, output_dim):
    inputs = torch.randn(num_samples, input_dim)
    targets = torch.randn(num_samples, output_dim)
    return inputs, targets

# 2. 定义简单的线性回归模型
class SimpleLinearModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(SimpleLinearModel, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.linear(x)

# 3. 定义训练函数
def train_model(model, dataloader, criterion, optimizer, scheduler, num_epochs, writer):
    model.train()
    for epoch in range(num_epochs):
        epoch_loss = 0.0
        for inputs, targets in dataloader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        avg_loss = epoch_loss / len(dataloader)
        writer.add_scalar('training_loss', avg_loss, epoch)
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}')
        
        # 动态调整学习率
        scheduler.step(avg_loss)
        writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], epoch)

# 4. 主函数
def main():
    # 参数设置
    num_samples = 100
    input_dim = 10
    output_dim = 1
    batch_size = 10
    num_epochs = 10
    learning_rate = 0.001

    # 生成随机数据
    inputs, targets = generate_random_data(num_samples, input_dim, output_dim)
    dataset = TensorDataset(inputs, targets)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # 初始化模型、损失函数和优化器
    model = SimpleLinearModel(input_dim, output_dim)
    criterion = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)

    # 初始化学习率调度器
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=2, verbose=True)

    # 初始化TensorBoard
    writer = SummaryWriter()

    # 训练模型
    train_model(model, dataloader, criterion, optimizer, scheduler, num_epochs, writer)

    # 关闭TensorBoard
    writer.close()

if __name__ == "__main__":
    main()